#include <stdio.h>
#include <string.h>
#include <string.h>
#include "hip/hip_runtime.h"
#include <hip/hip_ext.h>
#include <sys/time.h>

#define BLOCKDIM_X 512
#define BLOCKDIM_Y 1
#define BLOCKDIM_Z 1
#define GRIDDIM_X 1
#define GRIDDIM_Y 1
#define GRIDDIM_Z 1

#define TYPE_WIDTH 2

#define file_name "./kernel_test.co"
#define kernel_name "winograd_MiOpen_FP32_3x3S1x1_F2X3"

typedef struct DEMO_ARG_S
{
    int N;
    int C;
    int H;
    int W;
    int K;
    int n_groups;
    int flags;
    int reserved;
    float* in_ptr;
    float* weight_ptr;
    float* out_ptr;
    float* rsv_ptr = nullptr;
    int R;
    int S;
    int pad_h;
    int pad_w;
    int out_h;
    int out_w;
    float* bias_ptr = nullptr;
    float RELU_alpha;
    int reserved2;
    int64_t d_offset;
    int64_t f_offset;
    int64_t o_offset;
    int64_t b_offset;
    int d_N_stride;
    int d_C_stride;
    int d_H_stride;
    int d_W_stride;
    int f_K_stride;
    int f_C_stride;
    int f_R_stride;
    int f_S_stride;
    int o_N_stride;
    int o_K_stride;
    int o_H_stride;
    int o_W_stride;
    int G;
    int d_G_stride;
    int f_G_stride;
    int o_G_stride; 
}DEMO_ARG_S;

int main(int argc, char**argv)
{
    int batch = atoi(argv[1]);      // n
    int channels = atoi(argv[2]);   // c
    int height = atoi(argv[3]);   // h
    int width = atoi(argv[4]);    // w
    int k = atoi(argv[5]); // k
    int r = atoi(argv[6]);   // r
    int s = atoi(argv[7]);    // s
    int pad_h = atoi(argv[8]);
    int pad_w = atoi(argv[9]);
    int stride_h =atoi(argv[10]);
    int stride_w =atoi(argv[11]);
    hipSetDevice(atoi(argv[12]));
    int max_cu_num = 0;
    int device = 0;
    auto status = hipGetDevice(&device);
    hipDeviceGetAttribute(&max_cu_num, hipDeviceAttributeMultiprocessorCount, device);

    int o_w = floor(float(width + 2 * pad_w - s) / stride_w) + 1;
    int o_h = floor(float(height + 2 * pad_h - r) / stride_h) + 1;
    
    int kernel_total_item = k * channels * r * s;
    int input_total_item = batch * channels * height * width;
    int output_total_item = k * batch * o_h * o_w; 
    
    float* ori_kernel = (float*)malloc(kernel_total_item << TYPE_WIDTH);
    float* ori_input = (float*)malloc(input_total_item << TYPE_WIDTH);  
    
    float* input_device;
    float* kernel_device;
    float* output_device;
    hipMalloc((void**)&kernel_device, kernel_total_item << TYPE_WIDTH);
    hipMalloc((void**)&input_device, input_total_item << TYPE_WIDTH);
    hipMalloc((void**)&output_device, output_total_item << TYPE_WIDTH);

    //printf("ori kernel: ");
    for (int m = 0; m < k; m++) {
        for (int n = 0; n < channels; n++) {
            for (int o = 0; o < r; o++) {
                for (int p = 0; p < s; p++) {
                    int idx = m * channels * r * s + n * r * s + o * r + p;
                    ori_kernel[idx] =rand() % 16 / 1.0;
                    //ori_kernel[i] =rand() % 16 / 16.0;
                    //if (m == 0) {
                    //    printf("%.04f, ", ori_kernel[idx]);
                    //} 
                }   
            }   
        }   
    }   
    //printf("\n");
    //printf("ori input: ");
    for (int b = 0; b < batch; b++) {
        for (int c = 0; c < channels; c++) {
            for (int h = 0; h < height; h++) {
                for (int w = 0; w < width; w++) {
                    int idx = b * channels * height * width + c * height * width + h * width + w;
                    ori_input[idx] =rand() % 16 / 1.0;
                    //ori_input[i] =rand() % 16 / 16.0;
                    //if (b == 0 && (h >= 2 * 0 - pad_h && h < 2 * 0 - pad_h + 4) && (w >= 2 * 0 - pad_w && w < 2 * 0 - pad_w + 4)) {
                    //    printf("%d:%.04f, ", idx, ori_input[idx]);
                    //} 
                }   
            }   
        }
    }
    //printf("\n"); 

    hipMemcpy(kernel_device, ori_kernel, kernel_total_item << TYPE_WIDTH, hipMemcpyHostToDevice);
    hipMemcpy(input_device, ori_input, input_total_item << TYPE_WIDTH, hipMemcpyHostToDevice);

    constexpr int L_F_NKC_STRIDES   = 1 << 9;
    constexpr int L_F_GROUP_STRIDES = 1 << 10;
    DEMO_ARG_S kernel_arg;
    kernel_arg.N = batch;
    kernel_arg.C = channels;
    kernel_arg.H = height;
    kernel_arg.W = width;
    kernel_arg.K = k;
    kernel_arg.n_groups = max_cu_num;
    kernel_arg.flags = L_F_NKC_STRIDES + L_F_GROUP_STRIDES;
    kernel_arg.reserved = 0;
    kernel_arg.R = r;
    kernel_arg.S = s;
    kernel_arg.pad_h = pad_h;
    kernel_arg.pad_w = pad_w;
    kernel_arg.out_h = o_h;
    kernel_arg.out_w = o_w;
    kernel_arg.RELU_alpha = 0;
    kernel_arg.reserved2 = 0;
    kernel_arg.d_offset = 0;
    kernel_arg.f_offset = 0;
    kernel_arg.o_offset = 0;
    kernel_arg.b_offset = 0;
    kernel_arg.d_N_stride = channels * height * width * 4;
    kernel_arg.d_C_stride = height * width * 4;
    kernel_arg.d_H_stride = width * 4;
    kernel_arg.d_W_stride = 1 * 4;
    kernel_arg.f_K_stride = channels * r * s * 4;
    kernel_arg.f_C_stride = r * s * 4;
    kernel_arg.f_R_stride = s * 4;
    kernel_arg.f_S_stride = 1 * 4;
    kernel_arg.o_N_stride = k * o_h * o_w * 4;
    kernel_arg.o_K_stride = o_h * o_w * 4;
    kernel_arg.o_H_stride = o_w * 4;
    kernel_arg.o_W_stride = 1 * 4;
    kernel_arg.G = 1;
    kernel_arg.d_G_stride = 0;
    kernel_arg.f_G_stride = 0;
    kernel_arg.o_G_stride = 0;
    kernel_arg.in_ptr = input_device;
    kernel_arg.weight_ptr = kernel_device;
    kernel_arg.out_ptr = input_device;

    size_t param_size = sizeof(kernel_arg);

    void *config[] = {HIP_LAUNCH_PARAM_BUFFER_POINTER, &kernel_arg,
                        HIP_LAUNCH_PARAM_BUFFER_SIZE, &param_size, HIP_LAUNCH_PARAM_END};

    hipModule_t Module;
	hipFunction_t Function;

	hipModuleLoad(&Module,file_name);
	hipModuleGetFunction(&Function,Module,kernel_name);
    
    hipExtModuleLaunchKernel(Function, 
            max_cu_num * BLOCKDIM_X, 
            GRIDDIM_Y * BLOCKDIM_Y, 
            GRIDDIM_Z * BLOCKDIM_Z, 
            BLOCKDIM_X, 
            BLOCKDIM_Y, 
            BLOCKDIM_Z, 
            0, 0, NULL, (void **)&config, NULL, NULL, 0);
    float* output_cpu = (float*)malloc(output_total_item << TYPE_WIDTH);
    for (int i = 0; i < output_total_item; i++){
        output_cpu[i] = 0.0f;
    }
    hipMemcpy(output_cpu, output_device, output_total_item << TYPE_WIDTH, hipMemcpyDeviceToHost);

    hipEvent_t start, stop;
    hipEventCreate(&start);
    hipEventCreate(&stop);
    int run_times = 1;
    hipEventRecord(start, 0);
    for (int n = 0; n < run_times; n++) {
        hipExtModuleLaunchKernel(Function, 
        max_cu_num * BLOCKDIM_X, 
        GRIDDIM_Y * BLOCKDIM_Y, 
        GRIDDIM_Z * BLOCKDIM_Z, 
        BLOCKDIM_X, 
        BLOCKDIM_Y, 
        BLOCKDIM_Z, 
        0, 0, NULL, (void **)&config, NULL, NULL, 0);
    }
    hipEventRecord(stop, 0);
    hipEventSynchronize(stop);
    float time_elapsed_ms = 0;
    hipEventElapsedTime(&time_elapsed_ms, start, stop);
    printf("%f\n", time_elapsed_ms / run_times);
    hipEventDestroy(start);
    hipEventDestroy(stop);
    hipDeviceSynchronize();
    
    hipFree(kernel_device);
    hipFree(input_device);
    hipFree(output_device);

    free(ori_kernel);
    free(ori_input);
    free(output_cpu);

    return 0;
}
